import math

import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
from fvcore.nn.weight_init import c2_msra_fill, c2_xavier_fill
from .sparse_inst_loss import SparseInstCriterion, SparseInstMatcher

def _make_stack_3x3_convs(num_convs, in_channels, out_channels):
    convs = []
    for _ in range(num_convs):
        convs.append(
            nn.Conv2d(in_channels, out_channels, 3, padding=1))
        convs.append(nn.ReLU(True))
        in_channels = out_channels
    return nn.Sequential(*convs)


class MaskBranch(nn.Module):
    def __init__(self, cfg, in_channels):
        super().__init__()
        dim = cfg.hidden_dim
        num_convs = cfg.num_convs
        kernel_dim = cfg.kernel_dim
        self.mask_convs = _make_stack_3x3_convs(num_convs, in_channels, dim) # 创建多个卷积层，顺序为卷积 -> 激活 -> 卷积 -> 激活···
        self.projection = nn.Conv2d(dim, kernel_dim, kernel_size=1)
        self._init_weights()

    def _init_weights(self):
        for m in self.mask_convs.modules(): # 遍历self.mask_convs.modules()，也就是其中的卷积与激活
            if isinstance(m, nn.Conv2d): # 如果是卷积层，就初始化权重
                c2_msra_fill(m) # # MSRA初始化（Kaiming初始化的一种）
        c2_msra_fill(self.projection) # 对projection进行初始化，注意这里的projection层是卷积，但不是之前生成的mask_conv

    def forward(self, features):
        # mask features (x4 convs)
        features = self.mask_convs(features) # 使用生成的卷积层处理特征
        return self.projection(features) # 处理完之后再用一次卷积进行投影


class InstanceBranch(nn.Module):
    def __init__(self, cfg, in_channels, **kwargs):
        super().__init__()
        num_mask = cfg.num_query
        dim = cfg.hidden_dim
        num_classes = cfg.num_classes
        kernel_dim = cfg.kernel_dim
        num_convs = cfg.num_convs
        num_group = cfg.get('num_group', 1)
        sparse_num_group = cfg.get('sparse_num_group', 1)
        self.num_group = num_group
        self.sparse_num_group = sparse_num_group
        self.num_mask = num_mask
        self.inst_convs = _make_stack_3x3_convs(
                            num_convs=num_convs, 
                            in_channels=in_channels, 
                            out_channels=dim)

        self.iam_conv = nn.Conv2d(
            dim * num_group,
            num_group * num_mask * sparse_num_group,
            3, padding=1, groups=num_group * sparse_num_group)
        self.fc = nn.Linear(dim * sparse_num_group, dim)
        # output
        self.mask_kernel = nn.Linear(
            dim, kernel_dim)
        self.cls_score = nn.Linear(
            dim, num_classes)
        self.objectness = nn.Linear(
            dim, 1)
        self.prior_prob = 0.01
        self._init_weights()

    def _init_weights(self):
        for m in self.inst_convs.modules():
            if isinstance(m, nn.Conv2d):
                c2_msra_fill(m)
        bias_value = -math.log((1 - self.prior_prob) / self.prior_prob)
        for module in [self.iam_conv, self.cls_score]:
            init.constant_(module.bias, bias_value)
        init.normal_(self.iam_conv.weight, std=0.01)
        init.normal_(self.cls_score.weight, std=0.01)

        init.normal_(self.mask_kernel.weight, std=0.01)
        init.constant_(self.mask_kernel.bias, 0.0)
        c2_xavier_fill(self.fc)

    def forward(self, seg_features, is_training=True):
        out = {}
        # SparseInst part
        seg_features = self.inst_convs(seg_features) # 输入图像特征，然后卷积，对应Lane-level embedding的公式1
        # predict instance activation maps
        iam = self.iam_conv(seg_features.tile( # 将输入特征图 seg_features 沿通道维度复制 num_group 次
            (1, self.num_group, 1, 1)))
        if not is_training: # 训练的时候多用几组结果来找到最好参数，推理的时候要很多组结果没用，因为不知道哪个最好，所以只选一组，这样的话可以提高速度
            iam = iam.view(
                iam.shape[0],
                self.num_group,
                self.num_mask * self.sparse_num_group,
                *iam.shape[-2:])
            iam = iam[:, 0, ...]
            num_group = 1
        else:
            num_group = self.num_group

        iam_prob = torch.sigmoid(iam.clamp(min=-10, max=10))
        B, N = iam_prob.shape[:2]
        C = seg_features.size(1)
        # BxNxHxW -> BxNx(HW)
        iam_prob = iam_prob.view(B, N, -1)
        normalizer = iam_prob.sum(-1).clamp(min=1e-6)
        iam_prob_norm_hw = iam_prob / normalizer[:, :, None] # 将原始的实例激活图（iam）转换为归一化的概率分布，以便后续的特征聚合

        # aggregate features: BxCxHxW -> Bx(HW)xC
        # (B x N x HW) @ (B x HW x C) -> B x N x C
        all_inst_features = torch.bmm( # 批量矩阵乘法
            iam_prob_norm_hw,
            seg_features.view(B, C, -1).permute(0, 2, 1)) #BxNxC 通过激活图的权重加权求和，将空间特征转换为实例级特征，用于后续的分类、掩码生成等任务，对应公式2

        # concat sparse group features
        inst_features = all_inst_features.reshape(
            B, num_group,
            self.sparse_num_group,
            self.num_mask, -1
        ).permute(0, 1, 3, 2, 4).reshape(
            B, num_group,
            self.num_mask, -1)
        inst_features = F.relu_(
            self.fc(inst_features)) # 重组维度，然后激活，引入非线性并抑制负值

        # avg over sparse group
        iam_prob = iam_prob.view(
            B, num_group,
            self.sparse_num_group,
            self.num_mask,
            iam_prob.shape[-1]) # 将 [B, N, H*W] 转换为 [B, num_group, sparse_num_group, num_mask, H*W]
        iam_prob = iam_prob.mean(dim=2).flatten(1, 2) # 对每个实例的多个子组激活图取平均，提升预测稳定性、将 num_group 和 num_mask 维度展平，便于后续处理
        inst_features = inst_features.flatten(1, 2)
        out.update(dict(
            iam_prob=iam_prob,
            inst_features=inst_features)) # 得到结果，放进输出的字典中

        if self.training:
            pred_logits = self.cls_score(inst_features) # 将实例特征 C 维映射到 num_classes 维（如车道线类型：虚线、实线等）
            pred_kernel = self.mask_kernel(inst_features) # 动态卷积的权重参数，后面要与掩码分支的特征图做矩阵乘法，生成实例掩码
            pred_scores = self.objectness(inst_features) # 输出1维分数，表示实例是否为有效物体（前景/背景）
            out.update(dict(
                pred_logits=pred_logits,
                pred_kernel=pred_kernel,
                pred_scores=pred_scores)) # 将三类预测结果添加到 out 中，供后续损失计算使用
        return out

class SparseInsDecoder(nn.Module):
    def __init__(self, cfg, **kargs) -> None:
        super().__init__()
        in_channels = cfg.encoder.out_dims  # 移除 +2，不再拼接坐标特征
        self.output_iam = cfg.decoder.output_iam
        self.scale_factor = cfg.decoder.scale_factor
        self.sparse_decoder_weight = cfg.sparse_decoder_weight
        self.inst_branch = InstanceBranch(cfg.decoder, in_channels)
        self.mask_branch = MaskBranch(cfg.decoder, in_channels)
        self.sparse_inst_crit = SparseInstCriterion(
            num_classes=cfg.decoder.num_classes,
            matcher=SparseInstMatcher(),
            cfg=cfg)
        self._init_weights()

    def _init_weights(self):
        self.inst_branch._init_weights()
        self.mask_branch._init_weights()

    def forward(self, features, is_training=True, **kwargs):
        output = {}
        # 直接使用输入特征，不再拼接坐标特征
        inst_output = self.inst_branch(features, is_training=is_training)
        output.update(inst_output)

        if is_training:
            mask_features = self.mask_branch(features)
            pred_kernel = inst_output['pred_kernel']
            N = pred_kernel.shape[1]
            B, C, H, W = mask_features.shape

            pred_masks = torch.bmm(pred_kernel, mask_features.view(
                B, C, H * W)).view(B, N, H, W)
            pred_masks = F.interpolate(
                pred_masks, scale_factor=self.scale_factor,
                mode='bilinear', align_corners=False)
            output.update(dict(pred_masks=pred_masks))
        
        if self.training:
            sparse_inst_losses, matched_indices = self.loss(
                output,
                lane_idx_map=kwargs.get('lane_idx_map'),
                input_shape=kwargs.get('input_shape')
            )
            for k, v in sparse_inst_losses.items():
                sparse_inst_losses[k] = self.sparse_decoder_weight * v
            output.update(sparse_inst_losses)
            output['matched_indices'] = matched_indices
        return output

    def loss(self, output, lane_idx_map, input_shape):
        pred_masks = output['pred_masks']
        pred_masks = output['pred_masks'].view(
            pred_masks.shape[0],
            self.inst_branch.num_group,
            self.inst_branch.num_mask,
            *pred_masks.shape[2:])
        pred_logits = output['pred_logits']
        pred_logits = output['pred_logits'].view(
            pred_logits.shape[0],
            self.inst_branch.num_group,
            self.inst_branch.num_mask,
            *pred_logits.shape[2:])
        pred_scores = output['pred_scores']
        pred_scores = output['pred_scores'].view(
            pred_scores.shape[0],
            self.inst_branch.num_group,
            self.inst_branch.num_mask,
            *pred_scores.shape[2:])

        out = {}
        all_matched_indices = []
        for group_idx in range(self.inst_branch.num_group):
            sparse_inst_losses, matched_indices = \
                self.sparse_inst_crit(
                    outputs=dict(
                        pred_masks=pred_masks[:, group_idx, ...].contiguous(),
                        pred_logits=pred_logits[:, group_idx, ...].contiguous(),
                        pred_scores=pred_scores[:, group_idx, ...].contiguous(),
                    ),
                    targets=self.prepare_targets(lane_idx_map),
                    input_shape=input_shape,
                )
            for k, v in sparse_inst_losses.items():
                out['%s_%d' % (k, group_idx)] = v
            all_matched_indices.append(matched_indices)
        return out, all_matched_indices

    def prepare_targets(self, targets):
        new_targets = []
        for targets_per_image in targets:
            target = {}
            cls_labels = targets_per_image.flatten(-2).max(-1)[0]
            pos_mask = cls_labels > 0

            target["labels"] = cls_labels[pos_mask].long()
            target["masks"] = targets_per_image[pos_mask] > 0
            new_targets.append(target)
        return new_targets

class PositionalEncoding2D(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        # 将 2D 坐标映射到高维空间
        self.proj = nn.Sequential(
            nn.Conv2d(2, dim // 2, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(dim // 2, dim, kernel_size=1)
        )

    def forward(self, masks):
        """masks: [B, 2, H, W] (x_coords 和 y_coords)"""
        return self.proj(masks)  # [B, C, H, W]

class PolarEncoder(nn.Module):
    def __init__(self, in_channels, cfg=None, dim=256, window_size=5, stride=2):
        super().__init__()
        self.dim = dim
        self.in_channels = in_channels
        self.window_size = window_size
        self.stride = stride

        # If cfg is provided, override parameters from config
        if cfg is not None:
            self.dim = cfg.get('hidden_dim', dim)
            self.window_size = cfg.get('window_size', window_size)
            self.stride = cfg.get('stride', stride)

        # Feature extraction convs (from input channels to hidden dim)
        self.inst_convs = _make_stack_3x3_convs(
            num_convs=cfg.num_convs if cfg else 3,
            in_channels=in_channels,
            out_channels=self.dim
        )

        # Radial and angular projections
        self.r_proj = nn.Sequential(
            nn.Conv2d(1, self.dim // 2, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(self.dim // 2, self.dim // 2, kernel_size=1)
        )
        self.theta_proj = nn.Sequential(
            nn.Conv2d(1, self.dim // 2, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(self.dim // 2, self.dim // 2, kernel_size=1)
        )
        
        # Polar convolution
        self.vert_conv = nn.Conv2d(
            self.dim, self.dim,
            kernel_size=(self.window_size, 1),
            padding=(self.window_size // 2, 0),
            groups=1
        )

    def cartesian_to_polar(self, x_coords, y_coords):
        """将 2D 图像平面的 (x, y) 转换为极坐标 (r, θ)"""
        r = torch.sqrt(x_coords**2 + y_coords**2)  # 半径 [0, √2]
        theta = torch.atan2(y_coords, x_coords)    # 角度 [-π, π]
        return r, theta

    def sliding_window_scan(self, x):
        """沿极坐标角度 θ 方向滑动窗口（支持周期性填充）"""
        B, C, H, W = x.shape
        window_features = []
        
        for i in range(0, W, self.stride):
            start = i
            end = i + self.window_size
            if end > W:
                # 循环填充（角度 θ 是周期性的）
                pad = end - W
                window = torch.cat([x[:, :, :, start:], x[:, :, :, :pad]], dim=-1)
            else:
                window = x[:, :, :, start:end]
            
            # 卷积处理当前窗口
            window_feat = self.vert_conv(window)  # [B, C, H, 1]
            window_features.append(window_feat.squeeze(-1))  # [B, C, H]
        
        return torch.stack(window_features, dim=-1)  # [B, C, H, W']

    def forward(self, seg_features):
        """
        输入: seg_features [B, C, H, W] (2D 图像特征)
        输出: 极坐标编码后的特征 [B, dim, H, W']
        """
        B, C, H, W = seg_features.shape
        img_features = self.inst_convs(seg_features)
        
        # 1. 生成归一化 2D 坐标网格 (x, y) ∈ [-1, 1]
        y_coords = torch.linspace(-1, 1, H, device=seg_features.device).view(1, H, 1).expand(B, H, W)
        x_coords = torch.linspace(-1, 1, W, device=seg_features.device).view(1, 1, W).expand(B, H, W)
        
        # 2. 转换为极坐标 (r, θ)
        r, theta = self.cartesian_to_polar(x_coords, y_coords)
        polar_masks = torch.stack([r, theta], dim=1)  # [B, 2, H, W]
        
        # 3. 极坐标位置编码
        r_embed = self.r_proj(polar_masks[:, 0:1])    # [B, dim//2, H, W]
        theta_embed = self.theta_proj(polar_masks[:, 1:2])  # [B, dim//2, H, W]
        polar_embed = torch.cat([r_embed, theta_embed], dim=1)  # [B, dim, H, W]
        features = img_features + polar_embed
        
        # 4. 极坐标滑动窗口扫描
        window_features = self.sliding_window_scan(features)  # [B, dim, H, W']
        
        return window_features